#-*- conding: encoding -*-
#!/usr/bin/python3

'''
Author: ruanchao
Date: 2021-11-16 10:30:05
LastEditTime: 2021-12-03 08:32:44
LastEditors: ruanchao
Description: python 分析股票数据,来筛选一些合适买点的股票
FilePath: \stock-analysis-python\index.py
'''
import baostock as bs
import pandas as pd
from datetime import datetime
from time import time
from utils import get_trade_dates
import _thread
from mysqlHelper import mysqlHelper
from stockHelper import stockHelper
import sys

# 获取截止日期的数据
def get_end_date_stock_df(end_date):
    stock_rs = bs.query_all_stock(end_date)
    stock_df = stock_rs.get_data()
    if stock_df.empty == False:
        # 过滤股票代码
        # 1. 剔除停牌的股票
        # 2. 剔除指数
        # 3. 剔除bj.开头的股票
        # 4. 剔除68开头的股票
        # 5. 剔除code_name为空的股票
        stock_df = stock_df[stock_df['tradeStatus'] == '1']
        stock_df = stock_df[~stock_df['code'].str.startswith('sh.000')]
        stock_df = stock_df[~stock_df['code'].str.startswith('sz.399')]
        stock_df = stock_df[~stock_df['code'].str.startswith('bj.')]
        stock_df = stock_df[~stock_df['code'].str.startswith('sh.68')]
        stock_df = stock_df[stock_df['code_name'] != '']
    return stock_df

# 获取日期范围内的数据
def get_stock_df_data(dates):
    end_date = dates[len(dates) - 1]
    stock_df = get_end_date_stock_df(end_date)
    if stock_df.empty:
        print("日期{end_date}未获取到股票数据".format(end_date=end_date), '开始获取前一天的数据')
        end_date = dates[len(dates) - 2]
        print("前一天的日期为{end_date}".format(end_date=end_date))
        stock_df = get_end_date_stock_df(end_date)
    return [stock_df, end_date]

def add_trade_data_by_date(dates):
    # 获取指定日期的指数、股票数据
    start_date = dates[0]
    end_date = dates[len(dates) - 1]
    [stock_df, end_date] = get_stock_df_data(dates)
    if stock_df.empty:
        print("日期{end_date}和前一天未获取到股票数据".format(end_date=end_date))
        return False
    # 过滤一下st, 停牌
    for code in stock_df["code"]:
        print("Downloading :" + code + "开始")
        rs = bs.query_history_k_data_plus(code,"date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST",start_date, end_date, frequency="d", adjustflag="3")
        data_list = []
        while (rs.error_code == '0') & rs.next():
            # 获取一条记录，将记录合并在一起
            data_list.append(rs.get_row_data())
        # result = pd.DataFrame(data_list, columns=rs.fields)
        # print(result)
        #### 结果插入到数据库中 ####
        if (len(data_list)):
            # 每一个代码开启一个线程去处理
            # _thread.start_new_thread( add_stock_list_to_db, (data_list))
            add_stock_list_to_db(data_list)
        print("Downloading :" + code + "成功")
    return True

def add_stock_list_to_db(data_list):
    sql = "INSERT INTO `stock_list_baostock` (date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST) VALUES "
    for item in data_list:
        sql += "('{date}','{code}','{open}','{high}','{low}','{close}','{preclose}','{volume}','{amount}','{adjustflag}','{turn}','{tradestatus}','{pctChg}','{peTTM}','{pbMRQ}','{psTTM}','{pcfNcfTTM}','{isST}'),".format(date=item[0],code=item[1],open=item[2],high=item[3],low=item[4],close=item[5],preclose=item[6],volume=item[7],amount=item[8],adjustflag=item[9],turn=item[10],tradestatus=item[11],pctChg=item[12],peTTM=item[13],pbMRQ=item[14],psTTM=item[15],pcfNcfTTM=item[16],isST=item[17])
    db = mysqlHelper()
    sql = sql[:-1]
    db.execute_sql(sql)

def get_stock_list_data_from_db(sql):
    db = mysqlHelper()
    res = db.select(sql)
    return res

# 获取这个月需要进行插入数据的交易天数
def format_trade_dates(days):
    begin_date = days[0]
    end_date = days[len(days) - 1]
    begin_data = get_stock_list_data_from_db("select count(*) as count from `stock_list_baostock` where date='{}'".format(begin_date))
    end_data = get_stock_list_data_from_db("select count(*) as count from `stock_list_baostock` where date='{}'".format(end_date))
    begin_count = begin_data[0]["count"]
    end_count = end_data[0]["count"]
    print("begin_count", begin_count, end_count)
    # 如果第一天有数据,但是今天没有,并且是在17:30点之后,那么直接插入获取今天的数据
    if (begin_count > 0 and end_count == 0 and datetime.now().strftime("%H:%M") > "17:30"):
        return [end_date]
    # 如果第一天没数据，今天也没有，那么这个月没有进行更新，那么直接把这个月的数据都插入
    if (begin_count == 0 and end_count == 0):
        return days
    # 如果第一天和今天都有数据, 那么直接days就为空数组
    if (begin_count > 0 and end_count > 0):
        return []
    return []

def handle_data():
    start_time = time()
    print("获取股票数据开始时间: ", datetime.now())
    # 获取这个月的交易天数
    trade_days = get_trade_dates()
    print("交易天数: {trade_days}".format(trade_days=trade_days) )
    # 获取需要进行插入的交易天数
    format_trade_days = format_trade_dates(trade_days)
    if (len(format_trade_days) == 0):
        print("这个月没有需要插入数据的日期, 开始分析")
        return True
    # 遍历日期, 插入数据
    res = add_trade_data_by_date(format_trade_days)
    end_time = time()
    run_time = end_time - start_time
    print("获取股票数据结束时间: ", datetime.now(), ",耗时:", run_time)
    return res


#### 更新最后一天的股票基础数据
def update_last_day_stock_basic_data():
    # 获取到stock_list_baostock的code数据
    sql = "select code from stock_list_baostock where date = (select date from stock_list_baostock GROUP BY date ORDER BY date DESC limit 1)"
    stock_code_list = get_stock_list_data_from_db(sql)
    code_list = []
    print("stock_code_list", len(stock_code_list))
    if len(stock_code_list) > 0:
        for code_item in stock_code_list:
            code_list.append(code_item["code"])
    else:
        return
    # 获取stock_basic_baostock的数据
    sql = "select code from stock_basic_baostock"
    code_basic_list = get_stock_list_data_from_db(sql)
    exist_code_list = []
    if len(code_basic_list):
        for code_item in code_basic_list:
            exist_code_list.append(code_item["code"])
    for code in code_list:
        if (code in exist_code_list):
            # update_code_basic_data(code)
            continue
        else:
            insert_code_basic_data(code)
    return

#### 更新股票基本数据
def update_code_basic_data(code):
    code_data = get_stock_basic_data(code)
    if (len(code_data)):
        sql = "update stock_basic_baostock set code='{}',code_name='{}',ipoDate='{}',outDate='{}',type='{}',status='{}' where code='{}'".format(code_data["code"],code_data["code_name"],code_data["ipoDate"],code_data["outDate"],code_data["type"],code_data["status"],code_data["code"])
        print("update sql:", sql)
        db = mysqlHelper()
        db.execute_sql(sql)

def insert_code_basic_data(code):
    code_data = get_stock_basic_data(code)
    if (len(code_data)):
        sql = "insert into stock_basic_baostock (code, code_name, ipoDate, outDate, type, status) VALUES ('{}', '{}', '{}', '{}', '{}', '{}')".format(code_data["code"],code_data["code_name"],code_data["ipoDate"],code_data["outDate"],code_data["type"],code_data["status"])
        print("insert sql:", sql)
        db = mysqlHelper()
        db.execute_sql(sql)

def get_stock_basic_data(code):
    rs = bs.query_stock_basic(code)
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    code_dict = {}
    if len(data_list):
        columns = rs.fields
        result = pd.DataFrame(data_list, columns = rs.fields)
        for column in columns:
            code_dict[column] = result[column].values[0]
    return code_dict


############ 这里写股票分析算法 ###########
def analysis_data():
    # 获取金针探底的股票
    stockHelper().analysis_data()
    # 获取3日上涨的股票
    stockHelper().get_n_days_rise_stock(3)
    # 获取5日上涨的股票
    stockHelper().get_n_days_rise_stock(5)
    # 获取大涨之后回调的股票
    stockHelper().get_rise_callback_stock()
    # 获取某一天成交量放大的股票，并且成交量是前面N天平均值的X倍的票
    stockHelper().get_volume_rise_N_times_stock(10, 5)
    return



############ 这里写股票分析算法 ###########



if __name__ == '__main__':
    try:
        # 处理股票数据，在内部开启线程
        lg = bs.login()
        # 显示登陆返回信息
        print('login respond error_code:'+lg.error_code)
        print('login respond  error_msg:'+lg.error_msg)
        start_time = time()
        # 是否需要插入股票基础数据，建议每月更新一次
        update_base_data = False
        if update_base_data is True:
            update_last_day_stock_basic_data()
        # 处理数据
        data_res = handle_data()
        # 发现没有数据，那么去分析数据库中的数据
        analysis_data()
        end_time = time()
        run_time = end_time - start_time
        print("程序执行结束: ", datetime.now(), ",耗时:", run_time)
        bs.logout()
        sys.exit(1)
    except Exception as e:
        print ("Error: 程序出错:", e)
        sys.exit(1)

while 1:
   pass
